import os
from collections import defaultdict
from datetime import datetime
import numpy as np
import cv2
from ultralytics import YOLO
from scipy.spatial.distance import euclidean
import mysql.connector
import oss2  # 阿里云OSS SDK
import tensorflow as tf


# 创建OSS Bucket对象
auth = oss2.Auth(access_key_id, access_key_secret)
bucket = oss2.Bucket(auth, endpoint, bucket_name)

# 加载YOLOv8模型
model = YOLO('yolov10n.pt')

# 打开视频文件
video_path = r"C:\Users\H.muller\Desktop\交通大学路1.mp4"
video_filename = os.path.basename(video_path)
spot = video_filename.rsplit('.', 1)[0]  # 去掉文件扩展名
cap = cv2.VideoCapture(video_path)
assert cap.isOpened(), "Error reading video file"
fps = int(cap.get(cv2.CAP_PROP_FPS))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
duration = frame_count / fps

# 初始化像素到米的转换比例
PIXEL_TO_METER = 0.2

# 轨迹历史和车辆计数器
track_history = defaultdict(list)
total_vehicle_counter = 0
vehicle_counter = 0
start_time = 0
minute_vehicle_counts = []
total_minutes = 0

# 定义划线位置
LINE_POSITION = 250

# 轨迹平滑处理参数
SMOOTHING_WINDOW_SIZE = 5


# 定义类来存储车辆和行人的记录
class VehicleSpeedRecord:
    def __init__(self, id, label):
        self.id = id
        self.label = label
        self.speeds = []

    def add_speed(self, speed):
        self.speeds.append(speed)

    def calculate_average_speed(self):
        if len(self.speeds) == 0:
            return 0
        return sum(self.speeds) / len(self.speeds)


def get_video_duration(video_path):
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)

    # 读取帧率
    fps = cap.get(cv2.CAP_PROP_FPS)

    # 读取帧数
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 计算视频时长
    duration = frame_count / fps

    # 关闭视频文件
    cap.release()

    # 返回视频时长（以分钟为单位）
    return duration / 60


class PedestrianStateRecord:
    def __init__(self, id):
        self.id = id
        self.positions = []
        self.state = "unknown"

    def add_position(self, position):
        self.positions.append(position)

    def analyze_state(self):
        if len(self.positions) < 2:
            self.state = "standing"
        else:
            distances = [euclidean(self.positions[i], self.positions[i + 1]) for i in range(len(self.positions) - 1)]
            avg_distance = sum(distances) / len(distances)
            if avg_distance < 0.8:
                self.state = "standing"
            elif avg_distance < 10:
                self.state = "walking"
            else:
                self.state = "running"

    def __repr__(self):
        return f"PedestrianStateRecord(id={self.id}, positions={self.positions}, state={self.state})"


# 字典保存每个车辆和行人的记录
vehicle_speed_records = {}
pedestrian_state_records = {}

# 保存报警信息的列表
alert_records = []

# 记录每辆车上一次报警的时间
last_alert_time = defaultdict(lambda: datetime.min)


def calculate_speed(track_positions, fps, min_speed_threshold=0.5):
    if len(track_positions) < 2:
        return 0

    weighted_sum = 0.0
    total_time = 0.0

    for i in range(len(track_positions) - 1):
        current_pos = track_positions[i]
        next_pos = track_positions[i + 1]
        distance = euclidean(current_pos, next_pos)
        time_interval = 1 / fps

        # 计算未加权的速度
        speed_m_per_s_raw = distance * PIXEL_TO_METER / time_interval
        speed_km_per_h_raw = speed_m_per_s_raw * 3.6

        # 如果未加权速度低于阈值，将此段速度视为0
        if speed_km_per_h_raw < min_speed_threshold:
            continue

        center_y = (current_pos[1] + next_pos[1]) / 2
        offset = center_y - LINE_POSITION
        if center_y > LINE_POSITION:
            weight = 1 / (offset / 35 + 1)
        else:
            weight = 1 + (-offset) / (LINE_POSITION - 30)

        weighted_distance = distance * weight
        weighted_sum += weighted_distance
        total_time += time_interval

    if total_time == 0:
        return 0

    speed_m_per_s = weighted_sum * PIXEL_TO_METER / total_time
    speed_km_per_h = speed_m_per_s * 3.6

    return speed_km_per_h


def check_for_alerts(vehicle_id, speed, track_positions):
    reasons = []
    alert = False

    # 检查速度是否超过限制
    if speed > 70:
        alert = True
        reasons.append(f"车速过快 ({speed:.2f} km/h)")

    # 检查连续变道
    lane_change_count = 0
    if len(track_positions) >= SMOOTHING_WINDOW_SIZE:
        x_positions = [pos[0] for pos in track_positions[-SMOOTHING_WINDOW_SIZE:]]
        for i in range(1, len(x_positions)):
            if abs(x_positions[i] - x_positions[i - 1]) > 15:
                lane_change_count += 1
        if lane_change_count >= 3:
            alert = True
            reasons.append("连续变道")

    # 如果需要报警，记录报警信息
    alert_info = {}
    # 检查是否在3秒内重复报警
    current_time = datetime.now()
    time_since_last_alert = current_time - last_alert_time[vehicle_id]
    if alert and (time_since_last_alert.total_seconds() >= 3):
        # 仅记录每辆车每帧时间内的一次违规
        alert_info = {
            "vehicle_id": vehicle_id,
            "reasons": reasons,
            "speed": f"{speed:.2f}",
            "time": current_time.strftime('%Y-%m-%d %H:%M:%S')
        }
        alert_records.append(alert_info)
        last_alert_time[vehicle_id] = current_time

    return alert_info


# 添加记录车祸信息的类
class AccidentRecord:
    def __init__(self, vehicle_id, type, time):
        self.vehicle_id = vehicle_id
        self.type = type
        self.time = time

    def __repr__(self):
        return f"AccidentRecord(vehicle_id={self.vehicle_id}, type={self.type}, time={self.time})"


# 字典保存每个车辆的速度记录
vehicle_speed_records = defaultdict(list)

# 打开输出视频文件
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
output_path = r"C:\Users\H.muller\Desktop\processed.mp4"
out = cv2.VideoWriter(output_path, fourcc, fps,
                      (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

# 上传原视频到OSS
oss_original_path = f'videos/{video_filename}'
with open(video_path, 'rb') as f:
    bucket.put_object(oss_original_path, f, headers={'Content-Disposition': 'inline'})
print(f"Processed video uploaded to OSS: {oss_original_path}")

tracked_vehicles = set()  # 已经穿过线的车辆ID

while cap.isOpened():
    success, frame = cap.read()

    if success:
        results = model.track(frame, persist=True)
        speeds = []

        if results and len(results) > 0 and results[0].boxes is not None:
            boxes_xyxy = results[0].boxes.xyxy
            boxes_id = results[0].boxes.id
            boxes_cls = results[0].boxes.cls

            if boxes_xyxy is not None and boxes_id is not None and boxes_cls is not None:
                for i, (box, track_id, cls) in enumerate(zip(boxes_xyxy, boxes_id, boxes_cls)):
                    if track_id is None:
                        continue
                    center_x = (box[0] + box[2]) / 2
                    center_y = (box[1] + box[3]) / 2
                    position = (center_x, center_y)

                    track_history[track_id.item()].append(position)
                    label = model.model.names[int(cls.item())]

                    if label in ["car", "truck", "bus",]:
                        speed = calculate_speed(track_history[track_id.item()], fps)
                        # 忽略速度超过120 km/h的车辆
                        if speed > 120:
                            continue
                        speeds.append({
                            "id": int(track_id.item()),
                            "class": label,
                            "speed": speed
                        })

                        vehicle_id = int(track_id.item())
                        if vehicle_id not in vehicle_speed_records:
                            vehicle_speed_records[vehicle_id] = VehicleSpeedRecord(vehicle_id, label)
                        vehicle_speed_records[vehicle_id].add_speed(speed)

                        cv2.putText(frame, f"Speed: {speed:.2f} km/h",
                                    (int(center_x), int(center_y - 10)),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                        # 只在车辆首次穿过线时计数
                        if LINE_POSITION - 4 < position[1] < LINE_POSITION + 6 and vehicle_id not in tracked_vehicles:
                            vehicle_counter += 1
                            total_vehicle_counter += 1
                            tracked_vehicles.add(vehicle_id)

                        # 检查是否需要报警
                        alert_info = check_for_alerts(vehicle_id, speed, track_history[track_id.item()])
                        if alert_info:
                            print(f"警报！车辆ID: {vehicle_id} 原因: {', '.join(alert_info['reasons'])}")
                    elif label == "person":
                        pedestrian_id = int(track_id.item())
                        if pedestrian_id not in pedestrian_state_records:
                            pedestrian_state_records[pedestrian_id] = PedestrianStateRecord(pedestrian_id)
                        pedestrian_state_records[pedestrian_id].add_position(position)
                        pedestrian_state_records[pedestrian_id].analyze_state()
                        state = pedestrian_state_records[pedestrian_id].state

                        cv2.putText(frame, f"ID: {pedestrian_id} State: {state}",
                                    (int(center_x), int(center_y - 10)),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)

        print(speeds)

        annotated_frame = results[0].plot() if results and len(results) > 0 else frame
        cv2.line(annotated_frame, (0, LINE_POSITION), (annotated_frame.shape[1], LINE_POSITION), (0, 255, 0), 2)
        cv2.putText(annotated_frame, f"Traffic Flow: {total_vehicle_counter:.2f} vehicles",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        out.write(annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        break

# 获取视频总时长（分钟）
total_video_duration_minutes = get_video_duration(video_path)

# 计算平均车辆通过率
average_vehicles_per_minute = total_vehicle_counter / total_video_duration_minutes

# 计算并打印每个车辆的平均速度
for vehicle_id, record in vehicle_speed_records.items():
    average_speed = record.calculate_average_speed()
    print(f"Vehicle ID: {record.id}, Label: {record.label}, Average Speed: {average_speed:.2f} km/h")

# 输出行人状态记录
print("Pedestrian State Records:")
for record in pedestrian_state_records.values():
    print(record)

# 输出报警信息
print("报警信息：")
for record in alert_records:
    print(record)

# 将结果存入MySQL数据库
# MySQL连接参数
mysql_config = {
    'user': 'trafficSystem',  # 你的MySQL用户名
    'password': 'hhhhhh',  # 你的MySQL密码
    'host': '39.102.215.64',
    'database': 'trafficsystem',  # 要连接的数据库名
    'raise_on_warnings': True
}

# Upload processed video to Aliyun OSS
# Define OSS key and upload path
oss_key = 'processed_' + video_filename
oss_upload_path = f'videos/{oss_key}'

# 创建MySQL连接
cnx = mysql.connector.connect(**mysql_config)
cursor = cnx.cursor()

# 插入车辆速度记录到数据库
for vehicle_id, record in vehicle_speed_records.items():
    insert_speed_record = ("INSERT INTO vehicle "
                           "(label, average_speed, spot) "
                           "VALUES (%s, %s,%s)")
    average_speed = float(record.calculate_average_speed())
    print(average_speed)
    data_speed_record = (record.label, average_speed, spot)
    cursor.execute(insert_speed_record, data_speed_record)

# 插入平均每分钟车辆数到数据库
old_url = f'https://intelligent-traffic.oss-cn-beijing.aliyuncs.com/videos/{video_filename}'
new_url = f'https://intelligent-traffic.oss-cn-beijing.aliyuncs.com/videos/{oss_key}'

# 插入视频车辆数到数据库

insert_avg_per_minute = ("INSERT INTO flow "
                         "(old_url, new_url, time, average_flow, spot) "
                         "VALUES (%s, %s, %s, %s,%s)")
data_avg_per_minute = (old_url, new_url, datetime.now(), average_vehicles_per_minute, spot)
cursor.execute(insert_avg_per_minute, data_avg_per_minute)

# 插入报警信息
insert_alert = ("INSERT INTO warn "
                "(vehicle_id, reason, speed, time, spot) "
                "VALUES (%s,%s,%s,%s,%s)")
for record in alert_records:
    vehicle_id = record["vehicle_id"]
    reasons = "; ".join(record["reasons"])  # 将多个原因拼接成一个字符串
    speed = record["speed"]
    time = record["time"]
    data_alert = (vehicle_id, reasons, speed, time, spot)
    cursor.execute(insert_alert, data_alert)

# 提交事务
cnx.commit()

# 关闭游标和连接
cursor.close()
cnx.close()

# Release the video capture object and close display windows
cap.release()
out.release()
cv2.destroyAllWindows()

# 上传视频到oss云数据库
# with open(output_path, 'rb') as f:
#     bucket.put_object(oss_upload_path, f)
#
# print(f"Processed video uploaded to OSS: {oss_upload_path}")